import math
import torch
from torch import nn
from .attention_mask import AttentionMask
from .layer import EncoderLayer, DecoderLayer


class PositionalEncoding(nn.Module):
    """
    位置编码（Position Embedding）
    Transformer 中需要使用Position Embedding 表示单词出现在句子中的位置。
    因为 Transformer 不采用 RNN 的结构，而是使用全局信息，因此是无法捕捉到序列顺序信息的。
    例如将K、V按行进行打乱，那么Attention之后的结果是一样的。但是序列信息非常重要，代表着全局的结构，
    因此必须将序列的分词相对或者绝对position信息利用起来。
    """

    def __init__(self, d_model, context_length, dropout=0.1):
        """
        P_E(pos,2i) = sin(pos/(10000^(2i/d_model)))
        P_E(pos,2i+1) = cos(pos/(10000^(2i/d_model)))
        :param d_model: 单个词的特征维度
        :param context_length: 最大长度
        :param dropout:
        """
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(context_length, d_model)
        # shape: (context_length, d_model)

        # 生成一个 一维张量（tensor），包含从 0 开始到 context_length - 1 的连续浮点数序列
        position = torch.arange(0, context_length, dtype=torch.float)

        # 增加一个维度
        position = position.unsqueeze(1)
        # shape: (context_length, 1)

        # 1/10000^(2i/d_model) = exp(− 2i/d_model*log(10000.0) ) 利用指数函数e和对数函数log取下来，方便计算
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))

        # 偶数维度值
        pe[:, 0::2] = torch.sin(position * div_term)

        # 同理，奇数维度值
        pe[:, 1::2] = torch.cos(position * div_term)
        # shape: (context_length, d_model)

        batch_pe = pe.unsqueeze(0)
        # shape: (1, context_length, d_model)

        pe = batch_pe.transpose(0, 1)
        # shape: (context_length, 1, d_model)

        # 定一个缓冲区，其实简单理解为这个参数不更新就可以，但是参数仍然作为模型的参数保存
        self.register_buffer('pe', pe)

    def forward(self, x):
        """
        :param x: shape( context_length , batch_size, d_model)
        :return:
        """
        # 这里的self.pe是从缓冲区里拿的
        # 切片操作，把pe第一维的前context_length个tensor和x相加，其他维度不变
        # 实现词嵌入和位置编码的线性相加
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)


class Encoder(nn.Module):
    """
    Encoder部分包含三个部分：词向量embedding，位置编码部分，自注意力层及后续的前馈神经网络
    """

    def __init__(self, d_model, vocab_size, context_length, heads_number, n_layers, d_k, d_v, d_ff, dropout=0.1):
        super(Encoder, self).__init__()
        self.input_embedding = nn.Embedding(vocab_size, d_model)
        # 这行其实就是生成一个矩阵，算对每个单词编码 每个单词变成d_model=512维向量

        self.position_embedding = PositionalEncoding(d_model, context_length, dropout)
        # 位置编码，这里是固定的正余弦函数

        self.dropout_embedding = nn.Dropout(dropout)

        # 使用ModuleList对多个encoder进行堆叠
        self.layers = nn.ModuleList(
            [EncoderLayer(d_model, heads_number, d_k, d_v, d_ff, dropout) for _ in range(n_layers)])

        self.attention_mask = AttentionMask()

    def forward(self, encoder_inputs):
        # encoder_inputs shape(batch_size, context_length)
        encoder_input_embedding = self.input_embedding(encoder_inputs)
        # 词向量shape(batch_size, context_length ,d_model)
        # 交换前两个维度（第0维和第1维）
        encoder_input_embedding_change = encoder_input_embedding.transpose(0, 1)
        # 词向量shape(context_length, batch_size, d_model)

        encoder_outputs = self.position_embedding(encoder_input_embedding_change).transpose(0, 1)
        # 词向量shape(batch_size, context_length ,d_model)

        len_k = encoder_inputs.size()[1]
        padding_mask = self.attention_mask.get_padding_mask(encoder_inputs, len_k)
        encoder_self_attentions = []
        for layer in self.layers:
            # 送入6个mulit_head_attention
            encoder_outputs, encoder_self_attention = layer(encoder_outputs, padding_mask)
            encoder_self_attentions.append(encoder_self_attention)
        return encoder_outputs, encoder_self_attentions


class Decoder(nn.Module):
    """
    Decoder 部分包含三个部分：词向量embedding，位置编码部分，自注意力层及后续的前馈神经网络
    """

    def __init__(self, d_model, vocab_size, context_length, heads_number, n_layers, d_k, d_v, d_ff, dropout):
        """
        :param d_model: 嵌入向量的维度，它体现了每个单词会被表示成多少维的向量
        :param vocab_size: 目标语言词汇表的大小，也就是目标语言里不同单词的总数
        :param context_length: 单条信息的最大长度
        :param heads_number: 多头数量
        :param n_layers: 解码层数
        :param d_k:
        :param d_v:
        :param d_ff:
        :param dropout:
        """
        super(Decoder, self).__init__()
        self.input_embedding = nn.Embedding(vocab_size, d_model)
        self.position_embedding = PositionalEncoding(d_model, context_length, dropout)
        self.dropout_embedding = nn.Dropout(dropout)
        self.layers = nn.ModuleList(
            [DecoderLayer(d_model, heads_number, d_k, d_v, d_ff, dropout) for _ in range(n_layers)])
        self.attention_mask = AttentionMask()

    def forward(self, decoder_inputs, encoder_inputs, encoder_outputs):
        """
        :param decoder_inputs: shape( context_length , batch_size)
        :param encoder_inputs: shape( context_length , batch_size)
        :param encoder_outputs: shape( context_length , batch_size, d_model)
        """
        decoder_inputs_embedding = self.input_embedding(decoder_inputs)
        # 词向量shape(batch_size, max_len ,d_model)
        # 交换前两个维度（第0维和第1维）
        decoder_inputs_embedding_change = decoder_inputs_embedding.transpose(0, 1)
        # 词向量shape(max_len, batch_size, d_model)
        add_position_embedding = self.position_embedding(decoder_inputs_embedding_change).transpose(0, 1)
        # 词向量shape(batch_size, max_len, d_model)
        decoder_dropout_inputs = self.dropout_embedding(add_position_embedding)

        # get_padding_mask inputs自注意力层的时候的pad 部分
        decoder_len_k = decoder_inputs.size()[1]
        padding_mask = self.attention_mask.get_padding_mask(decoder_inputs, decoder_len_k)
        # shape(batch_size, max_len, max_len)

        # get_subsequent_mask 这个做的是自注意层的mask部分，就是当前单词之后看不到，使用一个上三角为1的矩阵
        subsequent_mask = self.attention_mask.get_subsequent_mask(decoder_inputs)
        # shape(batch_size, max_len, max_len)

        # 两个矩阵相加，大于0的为1，不大于0的为0(不是pad且没被掩码)，为1(是pad但掩码1+0，是pad且没掩码1+1)的在之后就会被fill到无限小  我们只喜欢False
        # inputs自注意力层的时候的pad部分 + inputs自注意层的mask部分
        self_attention_mask = torch.gt((padding_mask + subsequent_mask), 0)

        # 这个做的是交互注意力机制中的mask矩阵，enc的输入是K,V inputs的输入是Q  这里只关心k，不关心Q了
        encoder_len_k = encoder_inputs.size()[1]
        decoder_encoder_attention_mask = self.attention_mask.get_padding_mask(decoder_inputs, encoder_len_k)

        decoder_self_attentions, decoder_encoder_attentions = [], []
        decoder_outputs = None
        for layer in self.layers:
            layer_out = layer(decoder_dropout_inputs, encoder_outputs, self_attention_mask,
                              decoder_encoder_attention_mask)
            decoder_outputs, decoder_self_attention, decoder_encoder_attention = layer_out
            decoder_self_attentions.append(decoder_self_attention)
            decoder_encoder_attentions.append(decoder_encoder_attention)
        return decoder_outputs, decoder_self_attentions, decoder_encoder_attentions
